#
#  Copyright 2024 The InfiniFlow Authors. All Rights Reserved.
#
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.

# from beartype import BeartypeConf
# from beartype.claw import beartype_all  # <-- you didn't sign up for this
# beartype_all(conf=BeartypeConf(violation_type=UserWarning))    # <-- emit warnings from all code
import random
import sys
import threading
import time

from api.utils.log_utils import initRootLogger, get_project_base_directory
from graphrag.general.index import run_graphrag
from graphrag.utils import get_llm_cache, set_llm_cache, get_tags_from_cache, set_tags_to_cache
from rag.prompts import keyword_extraction, question_proposal, content_tagging

import logging
import os
from datetime import datetime
import json
import xxhash
import copy
import re
from functools import partial
from io import BytesIO
from multiprocessing.context import TimeoutError
from timeit import default_timer as timer
import tracemalloc
import signal
import trio
import exceptiongroup
import faulthandler

import numpy as np
from peewee import DoesNotExist

from api.db import LLMType, ParserType, TaskStatus
from api.db.services.document_service import DocumentService
from api.db.services.llm_service import LLMBundle
from api.db.services.task_service import TaskService
from api.db.services.file2document_service import File2DocumentService
from api import settings
from api.versions import get_ragflow_version
from api.db.db_models import close_connection
from rag.app import laws, paper, presentation, manual, qa, table, book, resume, picture, naive, one, audio, \
    email, tag
from rag.nlp import search, rag_tokenizer
from rag.raptor import RecursiveAbstractiveProcessing4TreeOrganizedRetrieval as Raptor
from rag.settings import DOC_MAXIMUM_SIZE, SVR_CONSUMER_GROUP_NAME, get_svr_queue_name, get_svr_queue_names, print_rag_settings, TAG_FLD, PAGERANK_FLD
from rag.utils import num_tokens_from_string, truncate
from rag.utils.redis_conn import REDIS_CONN, RedisDistributedLock
from rag.utils.storage_factory import STORAGE_IMPL
from graphrag.utils import chat_limiter

BATCH_SIZE = 64

# 定义解析器工厂映射表，将不同的文档类型映射到对应的解析函数或类
FACTORY = {
    "general": naive,
    ParserType.NAIVE.value: naive,
    ParserType.PAPER.value: paper,
    ParserType.BOOK.value: book,
    ParserType.PRESENTATION.value: presentation,
    ParserType.MANUAL.value: manual,
    ParserType.LAWS.value: laws,
    ParserType.QA.value: qa,
    ParserType.TABLE.value: table,
    ParserType.RESUME.value: resume,
    ParserType.PICTURE.value: picture,
    ParserType.ONE.value: one,
    ParserType.AUDIO.value: audio,
    ParserType.EMAIL.value: email,
    ParserType.KG.value: naive,
    ParserType.TAG.value: tag
}

# 用于存储未确认消息的迭代器，初始为 None，运行时动态赋值
UNACKED_ITERATOR = None

# 根据命令行参数设置消费者编号，默认为 "0"
CONSUMER_NO = "0" if len(sys.argv) < 2 else sys.argv[1]
# 构造消费者名称，格式为 task_executor_编号
CONSUMER_NAME = "task_executor_" + CONSUMER_NO
# 记录当前服务启动的时间戳，带时区信息，精确到毫秒
BOOT_AT = datetime.now().astimezone().isoformat(timespec="milliseconds")
# 初始化任务统计变量：待处理任务数、延迟任务数、已完成任务数、失败任务数
PENDING_TASKS = 0
LAG_TASKS = 0
DONE_TASKS = 0
FAILED_TASKS = 0

# 用于记录当前正在执行的任务字典，键为任务ID，值为任务对象
CURRENT_TASKS = {}

# 从环境变量中读取最大并发任务数，默认为5
MAX_CONCURRENT_TASKS = int(os.environ.get('MAX_CONCURRENT_TASKS', "5"))
# 从环境变量中读取最大并发分块构建器数量，默认为1
MAX_CONCURRENT_CHUNK_BUILDERS = int(os.environ.get('MAX_CONCURRENT_CHUNK_BUILDERS', "1"))
# 从环境变量中读取最大并发 MinIO 操作数，默认为10
MAX_CONCURRENT_MINIO = int(os.environ.get('MAX_CONCURRENT_MINIO', '10'))
# 创建 trio 的任务并发限制器，控制同时运行的任务数量
task_limiter = trio.CapacityLimiter(MAX_CONCURRENT_TASKS)
# 创建 trio 的分块处理并发限制器
chunk_limiter = trio.CapacityLimiter(MAX_CONCURRENT_CHUNK_BUILDERS)
# 创建 trio 的 MinIO 操作并发限制器
minio_limiter = trio.CapacityLimiter(MAX_CONCURRENT_MINIO)
# 从环境变量中读取心跳超时时间，默认为120秒
WORKER_HEARTBEAT_TIMEOUT = int(os.environ.get('WORKER_HEARTBEAT_TIMEOUT', '120'))
# 定义一个线程安全的停止事件对象，用于通知协程退出
stop_event = threading.Event()


# 定义信号处理函数 signal_handler，用于捕获并处理中断信号（如 Ctrl+C）
def signal_handler(sig, frame):
    # 记录接收到中断信号的日志信息，提示程序即将关闭
    logging.info("Received interrupt signal, shutting down...")
    # 设置全局的 stop_event 事件标志，通知其他线程或协程准备退出
    stop_event.set()
    # 等待1秒，确保其他线程有时间检测到 stop_event 并优雅退出
    time.sleep(1)
    # 终止当前进程，返回状态码0表示正常退出
    sys.exit(0)


# SIGUSR1 handler: start tracemalloc and take snapshot
# 定义一个信号处理函数，用于启动内存追踪并保存快照
def start_tracemalloc_and_snapshot(signum, frame):
    # 如果当前没有启用内存追踪，则开始 tracemalloc
    if not tracemalloc.is_tracing():
        logging.info("start tracemalloc")
        tracemalloc.start()
    else:
        logging.info("tracemalloc is already running")

    # 生成当前时间戳，用于命名快照文件
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    # 构建快照文件路径，包含进程 ID 和时间戳
    snapshot_file = f"snapshot_{timestamp}.trace"
    # snapshot_file的目录get_project_base_directory()-根目录ragflow 
    snapshot_file = os.path.abspath(os.path.join(get_project_base_directory(), "logs", f"{os.getpid()}_snapshot_{timestamp}.trace"))

    # 拍摄当前内存分配的快照，并将其写入文件
    snapshot = tracemalloc.take_snapshot()
    snapshot.dump(snapshot_file)
    # 获取当前和峰值内存使用量（仅跟踪部分）
    current, peak = tracemalloc.get_traced_memory()
    # 判断平台是否为 Windows，以选择不同的方式获取最大 RSS 内存
    if sys.platform == "win32":
        import  psutil
        process = psutil.Process()
        max_rss = process.memory_info().rss / 1024
    else:
        import resource
        # 获取当前进程的最大 RSS（resident set size）内存使用
        max_rss = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
    # 输出内存使用信息到日志中
    logging.info(f"taken snapshot {snapshot_file}. max RSS={max_rss / 1000:.2f} MB, current memory usage: {current / 10**6:.2f} MB, Peak memory usage: {peak / 10**6:.2f} MB")

# SIGUSR2 handler: stop tracemalloc
# 定义一个信号处理函数 stop_tracemalloc，用于在收到指定信号时停止内存追踪
def stop_tracemalloc(signum, frame):
    # 检查 tracemalloc（内存追踪模块）是否正在运行
    if tracemalloc.is_tracing():
        # 如果正在运行，则记录日志并停止内存追踪
        logging.info("stop tracemalloc")
        tracemalloc.stop()
    else:
        # 如果未运行，记录日志说明当前没有进行内存追踪
        logging.info("tracemalloc not running")

# 自定义异常类 TaskCanceledException，表示任务被取消的情况
class TaskCanceledException(Exception):
    # 构造函数接收一个消息参数，并将其保存为实例属性
    def __init__(self, msg):
        # 初始化异常消息内容
        self.msg = msg


# 定义 set_progress 函数，用于设置任务进度信息，支持分页和取消逻辑
def set_progress(task_id, from_page=0, to_page=-1, prog=None, msg="Processing..."):
    try:
        # 如果传入的进度值 prog 非空且小于0，则在消息前加上 [ERROR] 标记表示错误状态
        if prog is not None and prog < 0:
            msg = "[ERROR]" + msg
        # 查询当前任务是否已被标记为取消
        cancel = TaskService.do_cancel(task_id)

        # 如果任务被取消，在消息后追加 [Canceled]，并将进度设为 -1 表示中断
        if cancel:
            msg += " [Canceled]"
            prog = -1

        # 如果指定了 to_page 且大于0，说明是分页处理
        if to_page > 0:
            # 如果有消息内容，且 from_page 小于 to_page，则在消息前添加页码范围描述
            if msg:
                # 从开始页到结束页。
                if from_page < to_page:
                    msg = f"Page({from_page + 1}~{to_page + 1}): " + msg
        # 如果有消息内容，在前面加上当前时间戳
        if msg:
            msg = datetime.now().strftime("%H:%M:%S") + " " + msg
        # 构建进度更新字典，包含进度消息
        d = {"progress_msg": msg}
        # 如果进度值非空，也加入字典中
        if prog is not None:
            d["progress"] = prog

        # 调用 TaskService 更新任务进度信息
        TaskService.update_progress(task_id, d)

        # 关闭数据库连接以释放资源
        close_connection()
        # 如果任务被取消，抛出 TaskCanceledException 异常中断流程
        if cancel:
            raise TaskCanceledException(msg)
        # 记录本次进度更新的日志信息
        logging.info(f"set_progress({task_id}), progress: {prog}, progress_msg: {msg}")
    # 捕获 DoesNotExist 异常（如任务不存在），仅记录警告日志
    except DoesNotExist:
        logging.warning(f"set_progress({task_id}) got exception DoesNotExist")
    # 捕获所有其他异常，并记录完整堆栈信息以便排查问题
    except Exception:
        logging.exception(f"set_progress({task_id}), progress: {prog}, progress_msg: {msg}, got exception")

# 定义异步函数 collect，用于从 Redis Stream 中收集待处理的任务消息
async def collect():
    # 声明全局变量，用于记录消费者名称、已完成/失败任务数以及未确认的消息迭代器
    global CONSUMER_NAME, DONE_TASKS, FAILED_TASKS
    global UNACKED_ITERATOR
    # 获取当前服务关联的队列名称列表
    svr_queue_names = get_svr_queue_names()
    try:
        # 如果未初始化未确认消息迭代器，则通过 Redis 初始化它
        if not UNACKED_ITERATOR:
            UNACKED_ITERATOR = REDIS_CONN.get_unacked_iterator(svr_queue_names, SVR_CONSUMER_GROUP_NAME, CONSUMER_NAME)
        try:
            # 尝试从未确认消息迭代器中获取下一条消息
            redis_msg = next(UNACKED_ITERATOR)
        except StopIteration:
            # 如果迭代器已空（所有pending消息都已处理完），则尝试从队列中消费新消息
            for svr_queue_name in svr_queue_names:
                redis_msg = REDIS_CONN.queue_consumer(svr_queue_name, SVR_CONSUMER_GROUP_NAME, CONSUMER_NAME)
                if redis_msg:
                    break
    except Exception:
        # 捕获并记录任何异常，避免中断整个协程流程
        logging.exception("collect got exception")
        return None, None

    # 如果没有获取到消息，直接返回 None
    if not redis_msg:
        return None, None
    # 从消息对象中提取实际数据内容
    msg = redis_msg.get_message()
    # 如果消息内容为空，记录错误日志并确认该消息后返回
    if not msg:
        logging.error(f"collect got empty message of {redis_msg.get_msg_id()}")
        redis_msg.ack()
        return None, None

    # 初始化取消标志
    canceled = False
    # 根据消息中的任务ID查询任务信息
    task = TaskService.get_task(msg["id"])
    # 如果任务存在，进一步检查文档状态是否为取消或出错
    if task:
        _, doc = DocumentService.get_by_id(task["doc_id"])
        canceled = doc.run == TaskStatus.CANCEL.value or doc.progress < 0
    # 如果任务不存在或已被取消，增加失败计数，记录警告日志并确认消息
    if not task or canceled:
        # 状态标记
        state = "is unknown" if not task else "has been cancelled"
        FAILED_TASKS += 1
        logging.warning(f"collect task {msg['id']} {state}")
        redis_msg.ack()
        return None, None
    # 在任务数据中添加任务类型字段（如果消息中有提供）
    task["task_type"] = msg.get("task_type", "")
    # 返回原始消息对象和解析后的任务数据
    return redis_msg, task


# 定义异步函数 get_storage_binary，用于从指定存储桶中异步获取二进制数据
async def get_storage_binary(bucket, name):
    # 使用 trio.to_thread.run_sync 将同步的 STORAGE_IMPL.get 方法包装为异步调用
    return await trio.to_thread.run_sync(lambda: STORAGE_IMPL.get(bucket, name))


async def build_chunks(task, progress_callback):
    # 如果文件大小超过限制，设置进度为-1并返回空列表
    if task["size"] > DOC_MAXIMUM_SIZE:
        set_progress(task["id"], prog=-1, msg="File size exceeds( <= %dMb )" %
                                              (int(DOC_MAXIMUM_SIZE / 1024 / 1024)))
        return []

    # 根据parser_id选择对应的chunker
    chunker = FACTORY[task["parser_id"].lower()]
    try:
        # 记录开始时间
        st = timer()
        # 获取文档在存储中的地址
        bucket, name = File2DocumentService.get_storage_address(doc_id=task["doc_id"])
        # 从存储中获取二进制数据
        binary = await get_storage_binary(bucket, name)
        logging.info("From minio({}) {}/{}".format(timer() - st, task["location"], task["name"]))
    except TimeoutError:
        # 处理超时异常
        progress_callback(-1, "Internal server error: Fetch file from minio timeout. Could you try it again.")
        logging.exception(
            "Minio {}/{} got timeout: Fetch file from minio timeout.".format(task["location"], task["name"]))
        raise
    except Exception as e:
        # 处理其他异常
        if re.search("(No such file|not found)", str(e)):
            progress_callback(-1, "Can not find file <%s> from minio. Could you try it again?" % task["name"])
        else:
            progress_callback(-1, "Get file from minio: %s" % str(e).replace("'", ""))
        logging.exception("Chunking {}/{} got exception".format(task["location"], task["name"]))
        raise

    try:
        # 使用限流器进行异步分块处理
        """
        async with chunk_limiter:
            cks = await trio.to_thread.run_sync(
                lambda: chunker.chunk(
                    task["name"], 
                    binary=binary, 
                    from_page=task["from_page"],
                    to_page=task["to_page"], 
                    lang=task["language"], 
                    callback=progress_callback,
                    kb_id=task["kb_id"], 
                    parser_config=task["parser_config"], 
                    tenant_id=task["tenant_id"]
                )
            )
        
        """
        async with chunk_limiter:
            cks = await trio.to_thread.run_sync(lambda: chunker.chunk(task["name"], binary=binary, from_page=task["from_page"],
                                to_page=task["to_page"], lang=task["language"], callback=progress_callback,
                                kb_id=task["kb_id"], parser_config=task["parser_config"], tenant_id=task["tenant_id"]))
        logging.info("Chunking({}) {}/{} done".format(timer() - st, task["location"], task["name"]))
    except TaskCanceledException:
        # 处理任务取消异常
        raise
    except Exception as e:
        # 处理分块过程中的其他异常
        progress_callback(-1, "Internal server error while chunking: %s" % str(e).replace("'", ""))
        logging.exception("Chunking {}/{} got exception".format(task["location"], task["name"]))
        raise

    # 初始化docs列表和基础文档信息
    docs = []
    doc = {
        "doc_id": task["doc_id"],
        "kb_id": str(task["kb_id"])
    }
    # 如果有pagerank信息则添加
    if task["pagerank"]:
        doc[PAGERANK_FLD] = int(task["pagerank"])
    # 记录开始时间
    st = timer()

    # 定义上传单个块到MinIO的异步函数
    async def upload_to_minio(document, chunk):
        try:
            # 使用限流器上传图片
            # 因为考虑到并发，不要太大
            async with minio_limiter:
                # 深拷贝原始 document 字典，避免修改原始数据。
                d = copy.deepcopy(document)
                # 将 chunk 的键值对更新进新的字典 d，组合成最终的 chunk 文档数据
                d.update(chunk)
                # 生成唯一ID
                # 基于内容 + doc_id 使用了 xxhash 库的 xxh64 哈希算法，速度快且碰撞概率低。输入是 chunk["content_with_weight"] 和 doc_id 拼接后的字符串。
                # .hexdigest() 将哈希值转为十六进制字符串作为 ID。
                d["id"] = xxhash.xxh64((chunk["content_with_weight"] + str(d["doc_id"])).encode("utf-8")).hexdigest()
                # 添加创建时间和时间戳
                d["create_time"] = str(datetime.now()).replace("T", " ")[:19]
                # create_timestamp_flt：浮点型时间戳，表示从 epoch 开始经过的秒数。
                d["create_timestamp_flt"] = datetime.now().timestamp()
                # 如果没有图片则移除image字段
                if not d.get("image"):
                    _ = d.pop("image", None)
                    d["img_id"] = ""
                    docs.append(d)
                    return

                # 创建字节流缓冲区
                # 初始化一个内存中的字节流缓冲区，用于保存图像数据。
                output_buffer = BytesIO()
                if isinstance(d["image"], bytes):
                    # 如果 image 是 bytes 类型（如已压缩好的二进制图像），直接放入 BytesIO 缓冲区。
                    output_buffer = BytesIO(d["image"])
                else:
                    # 否则假设是 PIL 图像对象，调用 .save(...) 方法将其转换为 JPEG 格式。
                    d["image"].save(output_buffer, format='JPEG')
                # 上传到MinIO
                """
                await trio.to_thread.run_sync(
                    lambda: STORAGE_IMPL.put(
                        task["kb_id"], 
                        d["id"], 
                        output_buffer.getvalue()
                        )
                    )
                使用 trio.to_thread.run_sync(...) 包装，防止阻塞事件循环。

                """
                await trio.to_thread.run_sync(lambda: STORAGE_IMPL.put(task["kb_id"], d["id"], output_buffer.getvalue()))

                # 设置图片ID并移除image字段
                d["img_id"] = "{}-{}".format(task["kb_id"], d["id"])
                del d["image"]
                docs.append(d)
        except Exception:
            # 记录上传图片异常
            logging.exception(
                "Saving image of chunk {}/{}/{} got exception".format(task["location"], task["name"], d["id"]))
            raise

    # 并发执行上传到MinIO
    """
    trio.open_nursery()	创建一个“育儿所”（nursery），用来启动和管理多个子任务（协程）
    nursery.start_soon()	在 nursery 中异步启动一个新的任务（协程）
    nursery.start_soon(upload_to_minio, doc, ck) 并发上传，效率高
    nursery.start_soon(upload_to_minio, doc, ck, name=f"upload-{ck['id']}") 给每个任务起名
    """
    async with trio.open_nursery() as nursery:
        for ck in cks:
            nursery.start_soon(upload_to_minio, doc, ck)

    # 记录上传耗时
    el = timer() - st
    logging.info("MINIO PUT({}) cost {:.3f} s".format(task["name"], el))

    # 如果需要自动生成关键词
    if task["parser_config"].get("auto_keywords", 0):
        # 开始计时并更新进度
        st = timer()
        progress_callback(msg="Start to generate keywords for every chunk ...")
        # 加载语言模型
        chat_mdl = LLMBundle(task["tenant_id"], LLMType.CHAT, llm_name=task["llm_id"], lang=task["language"])

        # 定义生成关键词的异步函数
        async def doc_keyword_extraction(chat_mdl, d, topn):
            # 检查缓存
            cached = get_llm_cache(chat_mdl.llm_name, d["content_with_weight"], "keywords", {"topn": topn})
            if not cached:
                # 调用模型生成关键词
                async with chat_limiter:
                    # 通过keyword_extraction函数，可以看到调用大模型的prompt
                    cached = await trio.to_thread.run_sync(lambda: keyword_extraction(chat_mdl, d["content_with_weight"], topn))
                # 缓存结果
                set_llm_cache(chat_mdl.llm_name, d["content_with_weight"], cached, "keywords", {"topn": topn})
            if cached:
                # 存储关键词和token
                # 因为提示词让他用英文的逗号输出
                d["important_kwd"] = cached.split(",")
                d["important_tks"] = rag_tokenizer.tokenize(" ".join(d["important_kwd"]))
            return
        # 并发生成所有块的关键词
        async with trio.open_nursery() as nursery:
            for d in docs:
                nursery.start_soon(doc_keyword_extraction, chat_mdl, d, task["parser_config"]["auto_keywords"])
        # 更新进度并记录耗时
        progress_callback(msg="Keywords generation {} chunks completed in {:.2f}s".format(len(docs), timer() - st))

    # 如果需要自动生成问题
    if task["parser_config"].get("auto_questions", 0):
        # 开始计时并更新进度
        st = timer()
        progress_callback(msg="Start to generate questions for every chunk ...")
        # 加载语言模型
        chat_mdl = LLMBundle(task["tenant_id"], LLMType.CHAT, llm_name=task["llm_id"], lang=task["language"])

        # 定义生成问题的异步函数
        async def doc_question_proposal(chat_mdl, d, topn):
            # 检查缓存
            cached = get_llm_cache(chat_mdl.llm_name, d["content_with_weight"], "question", {"topn": topn})
            if not cached:
                # 调用模型生成问题
                async with chat_limiter:
                    cached = await trio.to_thread.run_sync(lambda: question_proposal(chat_mdl, d["content_with_weight"], topn))
                # 缓存结果
                set_llm_cache(chat_mdl.llm_name, d["content_with_weight"], cached, "question", {"topn": topn})
            if cached:
                # 存储问题和token
                d["question_kwd"] = cached.split("\n")
                d["question_tks"] = rag_tokenizer.tokenize("\n".join(d["question_kwd"]))
        # 并发生成所有块的问题
        async with trio.open_nursery() as nursery:
            for d in docs:
                nursery.start_soon(doc_question_proposal, chat_mdl, d, task["parser_config"]["auto_questions"])
        # 更新进度并记录耗时
        progress_callback(msg="Question generation {} chunks completed in {:.2f}s".format(len(docs), timer() - st))

    # 如果需要自动打标签
    if task["kb_parser_config"].get("tag_kb_ids", []):
        # 开始计时并更新进度
        progress_callback(msg="Start to tag for every chunk ...")
        # 获取相关参数
        kb_ids = task["kb_parser_config"]["tag_kb_ids"]
        tenant_id = task["tenant_id"]
        topn_tags = task["kb_parser_config"].get("topn_tags", 3)
        S = 1000
        # 开始计时
        st = timer()
        # 获取或生成标签
        examples = []
        all_tags = get_tags_from_cache(kb_ids)
        if not all_tags:
            all_tags = settings.retrievaler.all_tags_in_portion(tenant_id, kb_ids, S)
            set_tags_to_cache(kb_ids, all_tags)
        else:
            all_tags = json.loads(all_tags)

        # 加载语言模型
        chat_mdl = LLMBundle(task["tenant_id"], LLMType.CHAT, llm_name=task["llm_id"], lang=task["language"])

        # 准备需要打标签的文档
        docs_to_tag = []
        for d in docs:
            if settings.retrievaler.tag_content(tenant_id, kb_ids, d, all_tags, topn_tags=topn_tags, S=S) and len(d[TAG_FLD]) > 0:
                examples.append({"content": d["content_with_weight"], TAG_FLD: d[TAG_FLD]})
            else:
                docs_to_tag.append(d)

        # 定义打标签的异步函数
        async def doc_content_tagging(chat_mdl, d, topn_tags):
            # 检查缓存
            # 从缓存中尝试获取已有的标签结果，避免重复调用模型
            cached = get_llm_cache(chat_mdl.llm_name, d["content_with_weight"], all_tags, {"topn": topn_tags})
            if not cached:
                # 随机选取示例
                # 随机从示例列表中选取 2 个样本作为提示词中的 few-shot 示例
                picked_examples = random.choices(examples, k=2) if len(examples)>2 else examples
                if not picked_examples:
                    # 如果示例列表为空，则添加一个默认示例防止后续逻辑出错
                    picked_examples.append({"content": "This is an example", TAG_FLD: {'example': 1}})
                # 调用模型打标签
                # 使用限流器控制并发，防止超出模型 API 上限
                async with chat_limiter:
                    # 在线程池中同步调用内容打标签函数 content_tagging
                    cached = await trio.to_thread.run_sync(lambda: content_tagging(chat_mdl, d["content_with_weight"], all_tags, picked_examples, topn=topn_tags))
                # 如果成功获得模型返回结果，则将其转换为 JSON 字符串以便缓存存储
                if cached:
                    cached = json.dumps(cached)
            # 如果有缓存结果或刚刚生成了新结果
            if cached:
                # 缓存结果并更新文档
                # 将结果写入缓存，并更新文档对象中的标签字段
                set_llm_cache(chat_mdl.llm_name, d["content_with_weight"], cached, all_tags, {"topn": topn_tags})
                # 将缓存结果解析为字典并赋值给文档的 TAG_FLD 字段
                d[TAG_FLD] = json.loads(cached)
        # 并发打标签
        async with trio.open_nursery() as nursery:
            for d in docs_to_tag:
                nursery.start_soon(doc_content_tagging, chat_mdl, d, topn_tags)
        # 更新进度并记录耗时
        progress_callback(msg="Tagging {} chunks completed in {:.2f}s".format(len(docs), timer() - st))

    # 返回处理后的文档列表
    return docs


def init_kb(row, vector_size: int):
    # 根据租户 ID 构造对应的知识库索引名称
    idxnm = search.index_name(row["tenant_id"])
    # 调用文档存储连接对象，创建指定名称和向量维度的知识库索引
    return settings.docStoreConn.createIdx(idxnm, row.get("kb_id", ""), vector_size)


async def embedding(docs, mdl, parser_config=None, callback=None):
    # 如果 parser_config 为空，则初始化为空字典
    if parser_config is None:
        parser_config = {}
    # 设置默认批量大小为16，用于分批处理内容嵌入生成
    batch_size = 16
    # 初始化标题列表和内容列表
    tts, cnts = [], []
    # 遍历每个文档提取标题和内容，用于后续向量化
    for d in docs:
        # 获取文档标题，若无则使用默认值 "Title"
        tts.append(d.get("docnm_kwd", "Title"))
        # 尝试从问题字段获取内容，若为空则使用 content_with_weight 字段
        c = "\n".join(d.get("question_kwd", []))
        if not c:
            c = d["content_with_weight"]
        # 移除 HTML 表格标签以避免干扰嵌入模型输入
        c = re.sub(r"</?(table|td|caption|tr|th)( [^<>]{0,12})?>", " ", c)
        # 若内容仍为空，设置默认值 "None" 防止空输入
        if not c:
            c = "None"
        # 添加清洗后的内容到 cnts 列表
        cnts.append(c)

    # 初始化 token 数量计数器
    tk_count = 0
    # 如果标题数量与内容数量一致，对标题进行编码，并重复扩展以匹配内容维度
    if len(tts) == len(cnts):
        # 使用线程池运行模型 encode 方法处理标题部分
        vts, c = await trio.to_thread.run_sync(lambda: mdl.encode(tts[0: 1]))
        tts = np.concatenate([vts for _ in range(len(tts))], axis=0)
        tk_count += c

    # 初始化内容向量数组
    cnts_ = np.array([])
    # 分批次对内容进行编码，防止内存溢出并支持进度回调
    for i in range(0, len(cnts), batch_size):
        # 对当前批次的内容进行截断处理并调用模型编码
        vts, c = await trio.to_thread.run_sync(lambda: mdl.encode([truncate(c, mdl.max_length-10) for c in cnts[i: i + batch_size]]))
        # 将当前批次结果拼接到总结果中
        if len(cnts_) == 0:
            cnts_ = vts
        else:
            cnts_ = np.concatenate((cnts_, vts), axis=0)
        # 累加 token 计数
        tk_count += c
        # 调用回调函数更新进度信息
        callback(prog=0.7 + 0.2 * (i + 1) / len(cnts), msg="")
    # 将原始内容替换为编码后的向量数组
    cnts = cnts_

    # 从配置中读取标题向量权重，默认为 0.1
    title_w = float(parser_config.get("filename_embd_weight", 0.1))
    # 综合标题和内容向量，按照权重加权融合；若不一致则只使用内容向量
    vects = (title_w * tts + (1 - title_w) *
             cnts) if len(tts) == len(cnts) else cnts

    # 确保最终向量数量与输入文档数量一致
    assert len(vects) == len(docs)
    # 初始化向量维度变量
    vector_size = 0
    # 遍历每个文档，将对应的向量写入其字段中
    for i, d in enumerate(docs):
        v = vects[i].tolist()
        vector_size = len(v)
        d["q_%d_vec" % len(v)] = v
    # 返回 token 总数和向量维度
    return tk_count, vector_size


# 定义异步函数run_raptor，用于执行raptor处理流程
async def run_raptor(row, chat_mdl, embd_mdl, vector_size, callback=None):
    # 初始化空列表chunks用于存储文本片段和对应的向量
    chunks = []
    # 构造向量字段名称，格式为"q_{vector_size}_vec"
    vctr_nm = "q_%d_vec"%vector_size
    # 遍历检索器返回的chunk数据，提取内容和向量并添加到chunks中
    for d in settings.retrievaler.chunk_list(row["doc_id"], row["tenant_id"], [str(row["kb_id"])],
                                             fields=["content_with_weight", vctr_nm]):
        # 添加内容和向量（转换为numpy数组）到chunks列表
        chunks.append((d["content_with_weight"], np.array(d[vctr_nm])))

    # 创建Raptor实例，传入配置参数、模型、提示模板等信息
    raptor = Raptor(
        row["parser_config"]["raptor"].get("max_cluster", 64),
        chat_mdl,
        embd_mdl,
        row["parser_config"]["raptor"]["prompt"],
        row["parser_config"]["raptor"]["max_token"],
        row["parser_config"]["raptor"]["threshold"]
    )
    # 记录原始chunks长度，用于后续生成新内容时区分原始与摘要内容
    original_length = len(chunks)
    # 调用raptor实例进行处理，并等待异步结果
    chunks = await raptor(chunks, row["parser_config"]["raptor"]["random_seed"], callback)
    # 构建基础文档字典，包含文档ID、知识库ID、文档名和分词标题
    doc = {
        "doc_id": row["doc_id"],
        "kb_id": [str(row["kb_id"])],
        "docnm_kwd": row["name"],
        "title_tks": rag_tokenizer.tokenize(row["name"])
    }
    # 如果存在pagerank字段，则将其加入文档中
    if row["pagerank"]:
        doc[PAGERANK_FLD] = int(row["pagerank"])
    # 初始化结果列表res和token计数tk_count
    res = []
    tk_count = 0
    # 遍历生成的摘要内容及其向量
    for content, vctr in chunks[original_length:]:
        # 深拷贝基础文档结构到新文档d
        d = copy.deepcopy(doc)
        # 使用内容和doc_id生成唯一ID
        d["id"] = xxhash.xxh64((content + str(d["doc_id"])).encode("utf-8")).hexdigest()
        # 设置创建时间戳
        d["create_time"] = str(datetime.now()).replace("T", " ")[:19]
        # 设置浮点型时间戳
        d["create_timestamp_flt"] = datetime.now().timestamp()
        # 将向量转换为列表形式存入文档
        d[vctr_nm] = vctr.tolist()
        # 存储摘要内容
        d["content_with_weight"] = content
        # 对内容进行粗粒度分词
        d["content_ltks"] = rag_tokenizer.tokenize(content)
        # 对内容进行细粒度分词
        d["content_sm_ltks"] = rag_tokenizer.fine_grained_tokenize(d["content_ltks"])
        # 将构建好的文档添加到结果列表中
        res.append(d)
        # 累加当前内容的token数量
        tk_count += num_tokens_from_string(content)
    # 返回结果列表和总token数
    return res, tk_count


async def do_handle_task(task):
    # 从任务字典中提取任务 ID
    task_id = task["id"]
    # 提取分页起始页码
    task_from_page = task["from_page"]
    # 提取分页结束页码
    task_to_page = task["to_page"]
    # 提取租户 ID
    task_tenant_id = task["tenant_id"]
    # 提取嵌入模型 ID
    task_embedding_id = task["embd_id"]
    # 提取语言设置
    task_language = task["language"]
    # 提取 LLM 模型 ID
    task_llm_id = task["llm_id"]
    # 提取数据集 ID
    task_dataset_id = task["kb_id"]
    # 提取文档 ID
    task_doc_id = task["doc_id"]
    # 提取文档名称
    task_document_name = task["name"]
    # 提取解析器配置信息
    task_parser_config = task["parser_config"]
    # 记录任务开始时间戳
    task_start_ts = timer()

    # prepare the progress callback function
    # 创建进度回调函数，绑定当前任务的分页范围
    progress_callback = partial(set_progress, task_id, task_from_page, task_to_page)

    # FIXME: workaround, Infinity doesn't support table parsing method, this check is to notify user
    # 检查是否使用 Infinity 且选择了表格解析器，若满足条件则抛出异常并提示用户
    lower_case_doc_engine = settings.DOC_ENGINE.lower()
    if lower_case_doc_engine == 'infinity' and task['parser_id'].lower() == 'table':
        error_message = "Table parsing method is not supported by Infinity, please use other parsing methods or use Elasticsearch as the document engine."
        progress_callback(-1, msg=error_message)
        raise Exception(error_message)

    # 检查该任务是否已被取消
    task_canceled = TaskService.do_cancel(task_id)
    if task_canceled:
        # 如果任务被取消，则更新进度并直接返回
        progress_callback(-1, msg="Task has been canceled.")
        return

    try:
        # bind embedding model
        # 加载嵌入模型并进行简单测试编码，以验证模型可用性
        embedding_model = LLMBundle(task_tenant_id, LLMType.EMBEDDING, llm_name=task_embedding_id, lang=task_language)
        vts, _ = embedding_model.encode(["ok"])
        # 获取向量维度大小
        vector_size = len(vts[0])
    except Exception as e:
        # 若加载或调用模型失败，记录错误信息并抛出异常
        error_message = f'Fail to bind embedding model: {str(e)}'
        progress_callback(-1, msg=error_message)
        logging.exception(error_message)
        raise

    # 初始化知识库索引结构，准备后续向量写入
    init_kb(task, vector_size)

    # Either using RAPTOR or Standard chunking methods
    # 判断任务类型是否为 RAPTOR，决定是否启用递归树状聚类切片方法
    if task.get("task_type", "") == "raptor":
        # bind LLM for raptor
        # 绑定聊天模型用于 RAPTOR 处理中的推理环节
        chat_model = LLMBundle(task_tenant_id, LLMType.CHAT, llm_name=task_llm_id, lang=task_language)
        # run RAPTOR
        # 异步执行 RAPTOR 算法生成增强型文档块
        chunks, token_count = await run_raptor(task, chat_model, embedding_model, vector_size, progress_callback)
    # Either using graphrag or Standard chunking methods
    # 判断任务类型是否为 graphrag，决定是否启用知识图谱构建流程
    elif task.get("task_type", "") == "graphrag":
        # 声明并限制并发执行 graphrag 的协程数量上限
        global task_limiter
        task_limiter = trio.CapacityLimiter(2)
        # 如果配置中未启用 graphrag 功能，则直接返回
        if not task_parser_config.get("graphrag", {}).get("use_graphrag", False):
            return
        # 获取 graphrag 相关配置参数
        graphrag_conf = task["kb_parser_config"].get("graphrag", {})
        # 记录 graphrag 开始时间戳
        start_ts = timer()
        # 加载聊天模型用于图谱构建过程中的上下文推理
        chat_model = LLMBundle(task_tenant_id, LLMType.CHAT, llm_name=task_llm_id, lang=task_language)
        # 提取是否启用分辨率分析的布尔值
        with_resolution = graphrag_conf.get("resolution", False)
        # 提取是否启用社区发现机制的布尔值
        with_community = graphrag_conf.get("community", False)
        # 执行 graphrag 流程，构建文档级知识图谱
        await run_graphrag(task, task_language, with_resolution, with_community, chat_model, embedding_model, progress_callback)
        # 更新进度回调，表示知识图谱构建完成并记录耗时
        progress_callback(prog=1.0, msg="Knowledge Graph done ({:.2f}s)".format(timer() - start_ts))
        return
    else:
        # Standard chunking methods
        # 默认使用标准分块方式对文档内容进行切片处理
        start_ts = timer()
        # 异步构建文档块列表
        chunks = await build_chunks(task, progress_callback)
        # 输出日志记录文档构建耗时
        logging.info("Build document {}: {:.2f}s".format(task_document_name, timer() - start_ts))
        # 如果返回空 chunk 列表则直接返回
        if chunks is None:
            return
        # 若 chunk 列表为空，更新进度并提示无有效内容提取
        if not chunks:
            progress_callback(1., msg=f"No chunk built from {task_document_name}")
            return
        # TODO: exception handler
        ## set_progress(task["did"], -1, "ERROR: ")
        # 更新进度条信息，显示当前已生成的 chunk 数量
        progress_callback(msg="Generate {} chunks".format(len(chunks)))
        # 记录嵌入生成开始时间
        start_ts = timer()
        try:
            # 调用 embedding 函数为所有 chunk 生成向量表示
            token_count, vector_size = await embedding(chunks, embedding_model, task_parser_config, progress_callback)
        except Exception as e:
            # 捕获异常并更新进度状态为失败，同时记录错误日志
            error_message = "Generate embedding error:{}".format(str(e))
            progress_callback(-1, error_message)
            logging.exception(error_message)
            token_count = 0
            raise
        # 构建进度消息并输出至日志
        progress_message = "Embedding chunks ({:.2f}s)".format(timer() - start_ts)
        logging.info(progress_message)
        # 更新进度回调，显示嵌入生成阶段已完成
        progress_callback(msg=progress_message)

    # 统计唯一 chunk 数量，避免重复插入
    chunk_count = len(set([chunk["id"] for chunk in chunks]))
    # 记录文档写入开始时间戳
    start_ts = timer()
    # 初始化文档存储结果变量
    doc_store_result = ""
    # 设置每次批量插入 Elasticsearch/Infinity 的大小为 4
    es_bulk_size = 4
    # 分批次将 chunk 写入向量数据库或搜索引擎
    for b in range(0, len(chunks), es_bulk_size):
        # 使用线程池执行同步的插入操作，提升异步性能
        doc_store_result = await trio.to_thread.run_sync(lambda: settings.docStoreConn.insert(chunks[b:b + es_bulk_size], search.index_name(task_tenant_id), task_dataset_id))
        # 每隔 128 个 chunk 更新一次进度条，显示当前导入状态
        if b % 128 == 0:
            progress_callback(prog=0.8 + 0.1 * (b + 1) / len(chunks), msg="")
        # 如果插入失败，记录错误信息并抛出异常中断流程 插入操作，如果成功，返回的消息是None
        if doc_store_result:
            error_message = f"Insert chunk error: {doc_store_result}, please check log file and Elasticsearch/Infinity status!"
            progress_callback(-1, msg=error_message)
            raise Exception(error_message)
        # 提取已插入 chunk 的 ID 列表
        chunk_ids = [chunk["id"] for chunk in chunks[:b + es_bulk_size]]
        # 将 chunk ID 转换为空格分隔的字符串以便更新任务信息
        chunk_ids_str = " ".join(chunk_ids)
        try:
            # 更新任务中记录的 chunk ID 字符串，用于后续查询与管理
            TaskService.update_chunk_ids(task["id"], chunk_ids_str)
        except DoesNotExist:
            # 如果任务已被删除，则清理已插入的 chunk 数据并返回
            logging.warning(f"do_handle_task update_chunk_ids failed since task {task['id']} is unknown.")
            doc_store_result = await trio.to_thread.run_sync(lambda: settings.docStoreConn.delete({"id": chunk_ids}, search.index_name(task_tenant_id), task_dataset_id))
            return
    # 输出日志记录文档索引完成情况及耗时
    logging.info("Indexing doc({}), page({}-{}), chunks({}), elapsed: {:.2f}".format(task_document_name, task_from_page,
                                                                                     task_to_page, len(chunks),
                                                                                     timer() - start_ts))

    # 更新文档服务中的 chunk 数量、token 数量等统计信息
    DocumentService.increment_chunk_num(task_doc_id, task_dataset_id, token_count, chunk_count, 0)

    # 计算单个文档处理耗时和整个任务处理总耗时
    time_cost = timer() - start_ts
    task_time_cost = timer() - task_start_ts
    # 更新进度回调，表示所有阶段已完成，并展示各阶段耗时
    progress_callback(prog=1.0, msg="Indexing done ({:.2f}s). Task done ({:.2f}s)".format(time_cost, task_time_cost))
    # 输出最终任务完成日志，包含文档名、页码范围、chunk 数量、token 数量及耗时
    logging.info(
        "Chunk doc({}), page({}-{}), chunks({}), token({}), elapsed:{:.2f}".format(task_document_name, task_from_page,
                                                                                   task_to_page, len(chunks),
                                                                                   token_count, task_time_cost))


async def handle_task():
    # 声明全局变量 DONE_TASKS 和 FAILED_TASKS 用于统计任务状态
    global DONE_TASKS, FAILED_TASKS
    # 调用 collect 函数获取 Redis 消息和对应的任务数据
    redis_msg, task = await collect()
    # 如果没有获取到有效任务，则等待 5 秒后返回
    if not task:
        await trio.sleep(5)
        return
    try:
        # 记录任务开始日志，输出任务详情
        logging.info(f"handle_task begin for task {json.dumps(task)}")
        # 将当前任务加入 CURRENT_TASKS 字典，记录正在处理的任务
        CURRENT_TASKS[task["id"]] = copy.deepcopy(task)
        # 执行实际的任务处理逻辑
        await do_handle_task(task)
        # 增加已完成任务计数，并从 CURRENT_TASKS 中移除该任务
        DONE_TASKS += 1
        CURRENT_TASKS.pop(task["id"], None)
        # 记录任务完成日志
        logging.info(f"handle_task done for task {json.dumps(task)}")
    except Exception as e:
        # 捕获所有异常，增加失败任务计数，并从 CURRENT_TASKS 中移除该任务
        FAILED_TASKS += 1
        CURRENT_TASKS.pop(task["id"], None)
        try:
            # 提取错误信息，递归处理嵌套的异常组（如 Trio 的 ExceptionGroup）
            err_msg = str(e)
            while isinstance(e, exceptiongroup.ExceptionGroup):
                e = e.exceptions[0]
                err_msg += ' -- ' + str(e)
            # 设置任务进度为失败状态并附加错误信息
            set_progress(task["id"], prog=-1, msg=f"[Exception]: {err_msg}")
        except Exception:
            # 忽略设置进度时发生的任何异常
            pass
        # 记录异常详细堆栈信息
        logging.exception(f"handle_task got exception for task {json.dumps(task)}")
    # 确认消息已处理，向 Redis 发送 ACK
    redis_msg.ack()


async def report_status():
    # 声明多个全局变量，用于统计任务状态和上报心跳信息
    global CONSUMER_NAME, BOOT_AT, PENDING_TASKS, LAG_TASKS, DONE_TASKS, FAILED_TASKS
    # 将当前消费者名称加入 Redis 集合 TASKEXE，表示在线消费者
    REDIS_CONN.sadd("TASKEXE", CONSUMER_NAME)
    # 创建一个分布式 Redis 锁，用于清理过期消费者时的互斥操作
    redis_lock = RedisDistributedLock("clean_task_executor", lock_value=CONSUMER_NAME, timeout=60)
    # 进入无限循环，持续上报当前节点的状态和任务进度
    while True:
        try:
            # 获取当前时间戳，用于心跳上报的时间字段
            now = datetime.now()
            # 查询当前消费者组的任务队列状态信息
            group_info = REDIS_CONN.queue_info(get_svr_queue_name(0), SVR_CONSUMER_GROUP_NAME)
            # 如果获取到有效信息，则更新待处理任务数与滞后任务数
            if group_info is not None:
                PENDING_TASKS = int(group_info.get("pending", 0))
                LAG_TASKS = int(group_info.get("lag", 0))

            # 深拷贝当前正在处理的任务字典，防止并发写冲突
            current = copy.deepcopy(CURRENT_TASKS)
            # 构建心跳数据 JSON 字符串，包含当前节点的各种状态指标
            heartbeat = json.dumps({
                "name": CONSUMER_NAME,
                "now": now.astimezone().isoformat(timespec="milliseconds"),
                "boot_at": BOOT_AT,
                "pending": PENDING_TASKS,
                "lag": LAG_TASKS,
                "done": DONE_TASKS,
                "failed": FAILED_TASKS,
                "current": current,
            })
            # 将心跳数据写入 Redis 的有序集合中，以当前时间戳作为分值
            REDIS_CONN.zadd(CONSUMER_NAME, heartbeat, now.timestamp())
            # 记录本次心跳上报的日志信息
            logging.info(f"{CONSUMER_NAME} reported heartbeat: {heartbeat}")

            # 查询并删除超过30分钟未上报的心跳记录，避免数据堆积
            expired = REDIS_CONN.zcount(CONSUMER_NAME, 0, now.timestamp() - 60 * 30)
            if expired > 0:
                REDIS_CONN.zpopmin(CONSUMER_NAME, expired)

            # clean task executor
            # 尝试获取锁后进行消费者清理操作
            if redis_lock.acquire():
                # 获取所有已注册的消费者名称列表
                task_executors = REDIS_CONN.smembers("TASKEXE")
                for consumer_name in task_executors:
                    # 跳过当前消费者自身
                    if consumer_name == CONSUMER_NAME:
                        continue
                    # 查询该消费者最近是否有心跳上报，判断是否过期
                    expired = REDIS_CONN.zcount(
                        consumer_name, now.timestamp() - WORKER_HEARTBEAT_TIMEOUT, now.timestamp() + 10
                    )
                    # 若无心跳记录，则认为该消费者已下线，移除其注册信息
                    if expired == 0:
                        logging.info(f"{consumer_name} expired, removed")
                        REDIS_CONN.srem("TASKEXE", consumer_name)
                        REDIS_CONN.delete(consumer_name)
        except Exception:
            # 捕获并记录心跳上报过程中发生的任何异常
            logging.exception("report_status got exception")
        finally:
            # 确保无论如何都释放 Redis 分布式锁
            redis_lock.release()
        # 每隔30秒执行一次心跳上报和清理操作
        await trio.sleep(30)


# 创建一个 Redis 分布式锁，用于确保只有一个任务恢复进程在运行，防止重复处理
def recover_pending_tasks():
    # 初始化 Redis 分布式锁，使用唯一键名 "recover_pending_tasks"，设置当前消费者名称为锁值，超时时间为 60 秒
    redis_lock = RedisDistributedLock("recover_pending_tasks", lock_value=CONSUMER_NAME, timeout=60)
    # 获取所有服务队列名称列表，用于后续遍历检查每个队列中的待处理消息
    svr_queue_names = get_svr_queue_names()
    # 循环持续运行，直到 stop_event 被设置为止（例如接收到退出信号）
    while not stop_event.is_set():
        try:
            # 尝试获取分布式锁，成功后才进行后续操作，避免多个实例同时执行任务恢复
            if redis_lock.acquire():
                # 遍历所有服务队列名称
                for queue_name in svr_queue_names:
                    # 从 Redis 中获取指定队列中该消费组下的所有待处理消息
                    msgs = REDIS_CONN.get_pending_msg(queue=queue_name, group_name=SVR_CONSUMER_GROUP_NAME)
                    # 过滤掉已经被当前消费者处理的消息，只保留其他消费者未确认的消息
                    msgs = [msg for msg in msgs if msg['consumer'] != CONSUMER_NAME]
                    # 如果没有需要恢复的消息，跳过当前队列
                    if len(msgs) == 0:
                        continue

                    # 获取当前注册的所有任务执行器（TASKEXE）集合
                    task_executors = REDIS_CONN.smembers("TASKEXE")
                    # 构建任务执行器集合，用于判断消息所属消费者是否有效
                    task_executor_set = {t for t in task_executors}
                    # 进一步过滤消息：只保留那些其消费者不在当前活跃执行器列表中的消息
                    msgs = [msg for msg in msgs if msg['consumer'] not in task_executor_set]
                    # 遍历每一条待恢复的消息
                    for msg in msgs:
                        # 记录日志，显示正在恢复的消息 ID、原消费者名称以及消息已投递的时间
                        logging.info(
                            f"Recover pending task: {msg['message_id']}, consumer: {msg['consumer']}, "
                            f"time since delivered: {msg['time_since_delivered'] / 1000} s"
                        )
                        # 将消息重新入队，使其可以被重新消费处理
                        REDIS_CONN.requeue_msg(queue_name, SVR_CONSUMER_GROUP_NAME, msg['message_id'])
        except Exception:
            # 捕获并记录所有异常信息，提示任务恢复过程中出现错误
            logging.warning("recover_pending_tasks got exception")
        finally:
            # 无论是否出错，都释放之前获取到的分布式锁
            redis_lock.release()
            # 等待 60 秒后再次尝试恢复任务，形成周期性检测机制
            stop_event.wait(60)
        
async def task_manager():
    global task_limiter
    async with task_limiter:
        await handle_task()


# 定义异步主函数 main，用于启动和管理整个系统的运行流程
async def main():
    # 打印 ASCII 艺术风格的 RAGFlow 标志到日志中
    logging.info(r"""
  ______           __      ______                     __
 /_  __/___ ______/ /__   / ____/  _____  _______  __/ /_____  _____
  / / / __ `/ ___/ //_/  / __/ | |/_/ _ \/ ___/ / / / __/ __ \/ ___/
 / / / /_/ (__  ) ,<    / /____>  </  __/ /__/ /_/ / /_/ /_/ / /
/_/  \__,_/____/_/|_|  /_____/_/|_|\___/\___/\__,_/\__/\____/_/
    """)
    # 记录当前运行的 RAGFlow 版本号到日志中
    logging.info(f'TaskExecutor: RAGFlow version: {get_ragflow_version()}')
    # 初始化系统设置（从配置文件或环境变量加载）
    settings.init_settings()
    # 打印并记录 RAG 相关的配置信息（如文档大小限制、用户文件数量限制等）
    print_rag_settings()
    # 如果不是 Windows 平台，则注册内存跟踪信号处理函数 SIGUSR1 和 SIGUSR2
    if sys.platform != "win32":
        signal.signal(signal.SIGUSR1, start_tracemalloc_and_snapshot)
        signal.signal(signal.SIGUSR2, stop_tracemalloc)
    # 从环境变量读取是否启用内存分配追踪功能，默认为0（不启用）
    TRACE_MALLOC_ENABLED = int(os.environ.get('TRACE_MALLOC_ENABLED', "0"))
    # 如果启用了内存追踪，则主动调用开始追踪函数
    if TRACE_MALLOC_ENABLED:
        start_tracemalloc_and_snapshot(None, None)

    # 注册 Ctrl+C 或 kill 命令对应的信号处理函数，用于优雅关闭系统
    # 当接收到 SIGINT 信号（通常是用户按下 Ctrl+C）时，调用 signal_handler 函数进行处理
    signal.signal(signal.SIGINT, signal_handler)
    # 当接收到 SIGTERM 信号（通常是由其他程序发送，要求本程序终止）时，同样调用 signal_handler 函数进行处理
    signal.signal(signal.SIGTERM, signal_handler)

    # 启动一个守护线程用于恢复处于 pending 状态的任务
    threading.Thread(name="RecoverPendingTask", target=recover_pending_tasks).start()

    # 使用 trio 异步框架创建一个新的 nursery，用于并发执行协程任务
    async with trio.open_nursery() as nursery:
        # 在后台异步启动 report_status 协程，用于定期报告系统状态
        nursery.start_soon(report_status)
        # 循环检查 stop_event 是否被设置，未被设置时持续启动 task_manager 协程
        while not stop_event.is_set():
            # 每次循环都启动一个新的 task_manager 协程，用于处理任务队列
            nursery.start_soon(task_manager)
            # 防止过于频繁地启动任务，每次循环间隔 0.1 秒
            await trio.sleep(0.1)
    # 如果程序意外到达此处，输出严重错误日志提示开发者
    logging.error("BUG!!! You should not reach here!!!")

if __name__ == "__main__":
    faulthandler.enable()
    initRootLogger(CONSUMER_NAME)
    trio.run(main)
